# Licensed under the MIT license.

"""
rStar数学推理生成器 - 基于MCTS的多步推理系统

本模块实现了基于蒙特卡洛树搜索(MCTS)的数学推理生成系统，能够通过多种推理策略
(A1-A5)生成高质量的数学问题解答。

MCTS核心流程:
1. Selection(选择): 使用UCT公式选择最有潜力的节点
2. Expansion(扩展): 为选中节点生成子节点(通过A1-A5行为)
3. Simulation(模拟): 从扩展节点开始模拟到终端状态
4. Backpropagation(反向传播): 将奖励值向上传播更新节点统计

A1-A5推理行为详解:
- A1 (OST_STEP): 生成一步思考，逐步推进解题过程
- A2 (DIRECT_ANSWER): 直接生成最终答案，适用于简单问题
- A3 (SUBQUESTION): 将复杂问题分解为子问题，分而治之
- A4 (RE_SUBANSWER): 重新回答子问题，提供多样化解答
- A5 (REPHRASED_USER_QUESTION): 重新表述问题，从不同角度理解
"""

import sys
import os, json, time
from tqdm import tqdm

sys.path.append(".")

from common.utils import fix_seeds, setup_model_parallel, read_json
from common.arguments import get_parser, post_process_args, save_args
from run_src.rstar_utils import GeneratorError
from MCTS_for_reasoning import Generator, search_for_answers
from eval_src.Evaluator import *


def main(args):
    """
    主函数：执行基于MCTS的数学推理生成任务
    
    本函数实现完整的推理生成流程：
    1. 环境初始化 - 设置随机种子、模型并行等
    2. 数据加载 - 读取测试数据集和评估器
    3. 模型加载 - 根据API类型加载对应的语言模型
    4. 推理生成 - 对每个问题执行MCTS搜索生成解答
    5. 结果保存 - 保存生成的解答和统计信息

    参数:
        args (Namespace): 命令行参数对象，包含以下关键配置：
            - 模型配置: model_ckpt, api, tensor_parallel_size等
            - MCTS参数: num_rollouts, mcts_exploration_weight等  
            - 数据配置: dataset_name, data_root, test_json_filename等
            - 输出配置: outputs_dir, verbose等
    """
    # ================================ 环境初始化 ================================
    fix_seeds(args.seed)  # 固定随机种子确保实验可重现性
    
    if args.model_parallel:
        # 多卡并行模式：自动检测GPU数量并设置分布式环境
        args.local_rank, args.world_size = setup_model_parallel()
    else:
        # 单卡模式：使用默认设置
        args.local_rank, args.world_size = 0, 1

    # ================================ 数据加载 ================================
    # 构建测试数据文件完整路径
    test_file = os.path.join(args.data_root, args.dataset_name, args.test_json_filename + ".json")
    assert os.path.exists(test_file), f"测试文件不存在: {test_file}"
    
    # 读取测试数据列表，每个元素包含 {id, problem, solution}
    data_item_list = read_json(test_file)
    
    # 动态加载对应数据集的评估器 (如GSM8KEvaluator, MATHEvaluator等)
    evaluator = eval(f"{args.dataset_name}Evaluator()")

    # ================================ 模型加载 ================================
    tokenizer, model = None, None
    
    if args.api == "huggingface":
        # HuggingFace原生API：适用于本地部署的Transformer模型
        from models.HuggingFace_API import load_HF_model
        tokenizer, model = load_HF_model(args.model_ckpt)
        
    elif args.api == "vllm":
        # vLLM高性能推理引擎：支持张量并行和批量推理优化
        from models.vLLM_API import load_vLLM_model
        
        # 设置vLLM V1引擎专用环境变量，优化内存管理和性能
        os.environ['VLLM_DISABLE_STRICT_MEMORY_CHECK'] = '1'      # 绕过严格内存检查
        os.environ['VLLM_GPU_MEMORY_UTILIZATION'] = '0.85'       # GPU内存利用率85%
        os.environ['VLLM_TORCH_COMPILE_LEVEL'] = '0'             # 禁用torch编译优化
        os.environ['VLLM_DISABLE_CUSTOM_ALL_REDUCE'] = '1'       # 禁用自定义all-reduce
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'  # 启用可扩展内存段
        
        print("已设置vLLM V1引擎绕过配置")

        # 加载vLLM模型，支持多GPU张量并行
        tokenizer, model = load_vLLM_model(
            args.model_ckpt, 
            args.seed, 
            args.tensor_parallel_size,      # 张量并行度(如8卡并行)
            args.half_precision,            # 半精度推理
            args.max_num_seqs,             # 最大并发序列数
            args.qwen3_disable_thinking    # Qwen3特定优化
        )
        
    elif args.api == "gpt3.5-turbo":
        # OpenAI API：使用云端GPT模型
        from models.OpenAI_API import load_OpenAI_model
        tokenizer, model = load_OpenAI_model(args.model_ckpt)

    # 初始化推理生成器，集成模型、分词器和评估器
    generator = Generator(args, tokenizer, model, evaluator)

    # ================================ 统计变量初始化 ================================
    total_correct = 0           # 正确答案总数
    total_correct_limit = 0     # 正确答案限制数量
    num_tested = 0             # 已测试问题数量
    start_time = time.time()   # 记录总体开始时间

    # ================================ 主推理循环 ================================
    # 遍历所有测试问题，支持可视化进度条
    # disable逻辑: 只在主进程(local_rank=0)且verbose模式下显示进度条
    for i, data_item in enumerate(
        (pbar := tqdm(
            data_item_list, 
            disable=args.local_rank > 0 or not args.verbose, 
            position=1, 
            desc="处理问题"
        ))
    ):
        # 支持分片处理：只处理指定范围内的问题
        if i < args.start_idx or i >= args.end_idx:
            continue

        # 提取问题核心信息
        problem_id = data_item["id"]           # 问题唯一标识符
        problem = data_item["problem"]         # 问题文本描述
        gt_solution = data_item["solution"]    # 标准答案解题过程
        
        # 从标准答案中提取最终数值答案
        gt_answer = evaluator.extract_answer_from_gold_solution(gt_solution)

        # 在verbose模式下显示当前处理的问题详情
        if args.verbose:
            print(f"\n{'='*80}")
            print(f"正在处理问题 {i+1}/{len(data_item_list)} (ID: {problem_id})")
            print(f"问题: {problem[:200]}{'...' if len(problem) > 200 else ''}")
            print(f"正确答案: {gt_answer}")
            print(f"{'='*80}")

        # ================================ 结果数据结构初始化 ================================
        js = {
            "id": problem_id,                    # 问题ID
            "problem": problem,                  # 原始问题
            "model_completion": None,            # 模型生成的完整解答
            "model_answer": None,               # 模型提取的最终答案
            "all_model_completions": {},        # 所有生成的解答候选
            "gold_solution": gt_solution,       # 标准解题过程
            "gold_answer": gt_answer,           # 标准最终答案
        }

        # 初始化MCTS搜索结果容器
        model_solutions = []        # 每轮rollout的最佳解答
        stopping_id = -1           # 停止rollout的ID
        model_all_solutions = []   # 所有生成的解答方案

        # ================================ MCTS推理搜索 ================================
        """
        核心MCTS搜索过程：
        
        对当前问题执行num_rollouts轮MCTS搜索，每轮包含：
        1. Selection: 基于UCT值选择最有潜力的节点进行扩展
        2. Expansion: 通过A1-A5行为为选中节点生成多个子节点
           - A1: 生成一步思考(OST_STEP) - 逐步推进解题
           - A2: 生成直接答案(DIRECT_ANSWER) - 直接给出最终答案  
           - A3: 生成子问题(SUBQUESTION) - 分解复杂问题
           - A4: 重新回答子问题(RE_SUBANSWER) - 提供多样化解答
           - A5: 重新表述问题(REPHRASED_USER_QUESTION) - 换角度理解
        3. Simulation: 从扩展节点模拟到终端状态获得奖励
        4. Backpropagation: 将奖励值向上传播更新所有路径节点
        
        关键参数影响：
        - num_rollouts: 搜索轮数，越多越充分但耗时更长
        - mcts_exploration_weight: 探索权重，平衡探索vs利用
        - max_depth_allowed: 最大搜索深度，防止无限递归
        """
        start_problem_time = time.time()
        
        model_solutions, stopping_id, model_all_solutions = search_for_answers(
            args=args,
            user_question=problem,      # 待解决的数学问题
            question_id=i,             # 问题序号(用于日志和保存)
            gt_answer=gt_solution,      # 标准答案(用于评估)
            generator=generator         # 推理生成器实例
        )
        
        end_problem_time = time.time()
        
        # ================================ 结果数据填充 ================================
        # 从MCTS搜索结果中提取最佳解答
        if model_solutions and len(model_solutions) > 0:
            # 使用最后一轮的最佳解答作为模型的最终答案
            best_model_solution = model_solutions[-1]
            js["model_completion"] = best_model_solution
            
            # 从模型解答中提取最终答案
            if best_model_solution:
                js["model_answer"] = evaluator.extract_answer_from_model_completion(best_model_solution)
            else:
                js["model_answer"] = None
        else:
            js["model_completion"] = None
            js["model_answer"] = None
        
        # 填充所有候选解答
        js["all_model_completions"] = {
            f"rollout_{idx}": solution for idx, solution in enumerate(model_solutions)
        }
        
        # 添加MCTS搜索统计信息
        js["mcts_stats"] = {
            "num_rollouts": len(model_solutions),
            "stopping_rollout_id": stopping_id,
            "total_solutions_found": sum(len(solutions) for solutions in model_all_solutions),
            "processing_time_seconds": end_problem_time - start_problem_time,
            "api_calls_used": generator.io.call_counter,
            "tokens_used": generator.io.token_counter
        }
        
        # ================================ 结果统计与输出 ================================
        if args.verbose:
            print(f"\n问题 {i+1} 处理完成:")
            print(f"  - 用时: {end_problem_time - start_problem_time:.2f}秒")
            print(f"  - 生成解数量: {len(model_solutions)}")
            print(f"  - 最佳解答: {js['model_answer']}")
            print(f"  - 正确答案: {gt_answer}")
            print(f"  - 答案正确: {evaluator.check_answers_equiv(js['model_answer'], gt_answer) if js['model_answer'] else False}")
            print(f"  - 累计API调用: {generator.io.call_counter}")
            print(f"  - 累计Token使用: {generator.io.token_counter}")
        
        # 实时更新进度条描述，显示资源使用情况
        if args.verbose:
            pbar.set_description(
                f"处理问题 (API调用: {generator.io.call_counter}, "
                f"Token: {generator.io.token_counter})"
            )
        
        num_tested += 1  # 更新已测试问题计数

        # ================================ 结果保存 ================================
        # 保存单个问题的详细解答结果
        with open(os.path.join(args.answer_sheets_dir, f"Question {i:04d} - Answer.json"), "w") as f:
            json.dump(js, f, indent=2, ensure_ascii=False)

        # 保存中间统计结果，便于监控进度
        with open(os.path.join(args.run_outputs_dir, "intermediate_result.txt"), "w") as f:
            f.write(
                f"Total calls: {generator.io.call_counter}, "
                f"Avg calls: {generator.io.call_counter/(num_tested):.2f}\n"
            )
            f.write(
                f"Total tokens: {generator.io.token_counter}, "
                f"Avg tokens: {generator.io.token_counter/(num_tested):.2f}\n"
            )

    # ================================ 最终统计输出 ================================
    end_time = time.time()

    # 打印完整的性能统计信息
    print(f"==> 总API调用次数: {generator.io.call_counter}, "
          f"平均每题: {generator.io.call_counter/(num_tested):.2f}")
    print(f"==> 总Token使用量: {generator.io.token_counter}, "
          f"平均每题: {generator.io.token_counter/(num_tested):.2f}")
    print(f"==> 总耗时: {end_time-start_time:.2f}秒, "
          f"平均每题: {(end_time-start_time)/(num_tested):.2f}秒")

    # 保存最终完整统计结果
    with open(os.path.join(args.run_outputs_dir, "final_result.txt"), "w") as f:
        f.write(f"Total calls: {generator.io.call_counter}, "
                f"Avg calls: {generator.io.call_counter/(num_tested):.2f}\n")
        f.write(f"Total tokens: {generator.io.token_counter}, "
                f"Avg tokens: {generator.io.token_counter/(num_tested):.2f}\n")
        f.write(f"Total time: {end_time-start_time:.2f}s, "
                f"Avg time: {(end_time-start_time)/(num_tested):.2f}s\n")


if __name__ == "__main__":
    # ================================ 命令行参数配置 ================================
    """
    参数配置分为以下几个主要类别：
    
    1. 基础参数 (在common/arguments.py中定义)
       - model_ckpt: 模型检查点路径
       - api: 推理API类型 (huggingface/vllm/gpt3.5-turbo)
       - dataset_name: 数据集名称 (GSM8K/MATH/MULTIARITH等)
       - verbose: 详细输出控制
    
    2. MCTS核心参数
       - num_rollouts: MCTS搜索轮数，控制搜索充分程度
       - mcts_exploration_weight: UCT探索权重，平衡探索与利用
       - mcts_weight_scheduler: 权重调度策略 (const/exp/lin)
       - mcts_discount_factor: 奖励折扣因子，影响长期vs短期奖励
    
    3. 推理行为参数 (A1-A5)
       - num_a1_steps: A1行为生成的一步思考数量
       - num_subquestions: A3行为生成的子问题数量  
       - num_votes: 投票选择最佳答案的候选数量
       - max_depth_allowed: 最大搜索深度，防止过深递归
    
    4. 模型推理参数
       - tensor_parallel_size: 张量并行度 (多GPU加速)
       - max_num_seqs: 最大并发序列数 (批量推理)
       - half_precision: 是否使用半精度推理
    """
    parser = get_parser()  # 获取基础参数解析器

    # ================================ MCTS搜索参数 ================================
    parser.add_argument(
        "--num_rollouts", 
        type=int, 
        default=15,
        help="MCTS搜索轮数。每轮执行完整的Selection-Expansion-Simulation-Backpropagation过程。"
             "更多轮数提供更充分的搜索但增加计算成本。推荐范围: 10-50"
    )
    
    # ================================ 推理行为参数 (A1-A5) ================================
    """
    以下参数控制MCTS中A1-A5推理行为的具体实现机制：
    
    【核心实现原理】：
    这些参数并非简单地改变prompt，而是控制模型在每个MCTS节点扩展时生成多少个候选子节点。
    每个参数都会触发模型的多次推理调用，为同一个父节点创建多个不同的子节点，从而增加
    搜索树的分支因子和探索空间。
    
    【节点扩展机制】：
    当MCTS选中一个节点进行扩展时，会根据节点类型执行对应的A1-A5行为：
    - 每个行为会调用Generator中对应的generate_xxx方法
    - 这些方法内部调用IO_System.generate(num_return=参数值)
    - 模型会为相同的prompt生成num_return个不同的输出
    - 每个输出都会创建一个新的子节点并加入MCTS树
    
    【搜索树结构影响】：
    参数值直接决定了搜索树的宽度：
    - 更大的参数值 → 更多的子节点 → 更宽的搜索树 → 更丰富的探索
    - 但同时也意味着更多的API调用和计算开销
    """
    
    # ================================ A1行为：一步思考 (OST_STEP) ================================
    parser.add_argument(
        "--num_a1_steps", 
        type=int, 
        default=None,
        help="A1行为(一步思考)生成的候选步骤数量。"
             "\n【实现机制】："
             "\n  - 调用Generator.generate_ost_step()方法"
             "\n  - 内部执行: IO_System.generate(num_return=num_a1_steps)"
             "\n  - 模型为相同的思考提示生成num_a1_steps个不同的推理步骤"
             "\n  - 每个步骤创建一个OST_STEP类型的子节点"
             "\n【节点扩展】："
             "\n  - 父节点 → num_a1_steps个OST_STEP子节点"
             "\n  - 每个子节点包含一个独特的推理步骤文本"
             "\n  - 子节点可以继续向下扩展(递归A1-A2行为)"
             "\n【使用场景】："
             "\n  - 逐步推理：将复杂问题分解为多个思考步骤"
             "\n  - 多样化探索：同一状态下尝试不同的推理方向"
             "\n  - 渐进式求解：避免直接跳到最终答案"
             "\n【推荐范围】：2-5，默认3。更多候选提供更丰富的推理路径但增加计算成本"
    )
    
    parser.add_argument(
        "--disable_a1", 
        action="store_true",
        help="禁用A1行为(一步思考)。"
             "\n【影响】：禁用后节点将无法生成OST_STEP类型的子节点"
             "\n【适用场景】：简单问题或希望快速得到直接答案时"
             "\n【注意】：可能影响复杂问题的逐步推理能力"
    )

    # ================================ A3行为：子问题分解 (SUBQUESTION) ================================
    parser.add_argument(
        "--num_subquestions", 
        type=int, 
        default=3,
        help="A3行为生成的子问题候选数量。"
             "\n【实现机制】："
             "\n  - 调用Generator.generate_subquestions()方法"
             "\n  - 内部执行: IO_System.generate(num_return=num_subquestions)"
             "\n  - 模型为相同的分解提示生成num_subquestions个不同的子问题"
             "\n  - 每个子问题自动生成对应的子答案"
             "\n  - 每个(子问题,子答案)对创建一个SUBQUESTION类型的子节点"
             "\n【节点扩展】："
             "\n  - 父节点 → num_subquestions个SUBQUESTION子节点"
             "\n  - 每个子节点包含独特的问题分解策略"
             "\n  - 子节点可以继续分解或直接求解"
             "\n【分解策略】："
             "\n  - 将复杂问题拆分为更简单的子问题"
             "\n  - 不同候选提供不同的分解角度和粒度"
             "\n  - 支持多层次递归分解"
             "\n【推荐范围】：2-5，默认3。平衡分解多样性和计算效率"
    )

    # ================================ 投票选择参数 ================================
    parser.add_argument(
        "--num_votes", 
        type=int, 
        default=10,
        help="投票选择最佳答案时的候选数量。"
             "\n【实现机制】："
             "\n  - 在A3、A4行为中用于生成子答案候选"
             "\n  - 调用IO_System.generate(num_return=num_votes)"
             "\n  - 模型为子问题生成num_votes个不同答案"
             "\n  - 通过多数投票或概率评估选择最佳答案"
             "\n【质量保证】："
             "\n  - 更多候选 → 更可靠的答案选择"
             "\n  - 减少单次生成的随机性影响"
             "\n  - 提高答案的一致性和准确性"
             "\n【计算开销】："
             "\n  - 每个子问题都会生成num_votes个候选答案"
             "\n  - 总API调用次数 = num_subquestions × num_votes"
             "\n【推荐范围】：5-20，默认10。平衡答案质量和计算成本"
    )

    # ================================ 搜索深度控制 ================================
    parser.add_argument(
        "--max_depth_allowed", 
        type=int, 
        default=5,
        help="MCTS树的最大搜索深度限制。"
             "\n【实现机制】："
             "\n  - 在节点扩展时检查当前深度"
             "\n  - 超过限制的节点将停止生成子节点"
             "\n  - 防止无限递归和过深搜索"
             "\n【深度计算】："
             "\n  - 根节点(USER_QUESTION): depth=0"
             "\n  - 重述问题(REPHRASED_USER_QUESTION): depth=1"
             "\n  - 子问题/思考步骤: depth递增"
             "\n【性能影响】："
             "\n  - 过小：可能无法充分探索复杂问题"
             "\n  - 过大：计算开销大，可能陷入过细分解"
             "\n【推荐范围】：3-8，默认5。根据问题复杂度调整"
    )

    # ================================ MCTS算法核心参数 ================================
    parser.add_argument(
        "--mcts_discount_factor", 
        type=float, 
        default=1.0,
        help="MCTS奖励折扣因子。控制对未来奖励的重视程度。"
             "1.0表示不折扣(重视长期奖励)，<1.0表示偏向即时奖励。推荐范围: 0.8-1.0"
    )
    
    parser.add_argument(
        "--mcts_exploration_weight", 
        type=float, 
        default=2.0,
        help="UCT公式中的探索权重参数。平衡探索(exploration)与利用(exploitation)。"
             "较大值鼓励探索未访问节点，较小值偏向利用已知好节点。推荐范围: 1.0-3.0"
    )
    
    parser.add_argument(
        "--mcts_weight_scheduler", 
        choices=["exp", "lin", "const"], 
        default="const",
        help="探索权重的调度策略："
             "const-恒定权重，exp-指数衰减，lin-线性衰减。"
             "动态调度可以实现从探索到利用的平滑过渡"
    )
    
    parser.add_argument(
        "--mcts_num_last_votes", 
        type=int, 
        default=None,
        help="最终投票阶段的候选答案数量。用于从多个解答中选择最终答案。"
             "None表示使用默认值32。推荐范围: 16-64"
    )
    
    parser.add_argument(
        "--save_tree", 
        action="store_true",
        help="是否保存MCTS搜索树结构到文件。用于调试和分析搜索过程。"
             "启用会增加I/O开销但便于理解算法行为"
    )

    # ================================ 问题重述参数 (A5相关) ================================
    parser.add_argument(
        "--modify_prompts_for_rephrasing", 
        action="store_true",
        help="为重述任务修改提示词。启用后使用专门优化的重述提示模板，提高重述质量"
    )
    
    parser.add_argument(
        "--disable_a5", 
        action="store_true",
        help="禁用A5行为(问题重述)。A5通过重新表述问题提供不同理解角度，禁用会减少问题理解的多样性"
    )

    # ================================ 答案选择策略参数 ================================
    # enable_potential_score参数已在common/arguments.py中定义
    # 用于启用基于潜在分数的答案选择机制，提高答案质量评估的准确性

    # ================================ 参数后处理 ================================
    args = parser.parse_args()  # 解析命令行参数

    # 设置默认值：MCTS最终投票候选数量
    if args.mcts_num_last_votes is None:
        args.mcts_num_last_votes = 32  # 默认32个候选进行最终投票

    # 设置默认值：A1行为步骤数量
    if not args.disable_a1:
        if args.num_a1_steps is None:
            args.num_a1_steps = 3  # 默认生成3个一步思考候选

    # ================================ 提示词路径配置 ================================
    """
    根据数据集类型动态配置提示词文件路径。
    不同数据集(GSM8K/MATH/MULTIARITH等)使用不同的提示词模板以适应其特点：
    
    - fewshot_cot: Few-shot思维链提示，提供推理示例
    - fewshot_ost: Few-shot一步思考提示，指导逐步推理  
    - decompose: 问题分解提示，指导子问题生成
    - rephrasing: 问题重述提示，指导问题重新表达
    """
    prompts_dir = os.path.join(args.prompts_root, args.dataset_name)

    # Few-shot思维链提示配置
    args.fewshot_cot_prompt_path = os.path.join(
        prompts_dir, "fewshot_cot", "fewshot_cot_prompt.txt"
    )
    args.fewshot_cot_config_path = os.path.join(
        prompts_dir, "fewshot_cot", "fewshot_cot_config.json"
    )

    # Few-shot一步思考提示配置  
    args.fewshot_ost_prompt_path = os.path.join(
        prompts_dir, "fewshot_ost", "fewshot_ost_prompt.txt"
    )
    args.fewshot_ost_config_path = os.path.join(
        prompts_dir, "fewshot_ost", "fewshot_ost_config.json"
    )

    # 问题分解提示配置
    args.decompose_template_path = os.path.join(
        prompts_dir, "decompose", "decompose_template.json"
    )
    args.decompose_prompt_path = os.path.join(
        prompts_dir, "decompose", "decompose_prompt.txt"
    )

    # 问题重述相关配置 (仅在启用A5时需要)
    if not args.disable_a5:
        args.rephrasing_prompt_template_path = os.path.join(
            prompts_dir, "rephrasing_prompt_template.txt"
        )
        
        if args.modify_prompts_for_rephrasing:
            # 使用专门的重述优化提示词
            args.fewshot_cot_prompt_rephrased_path = os.path.join(
                prompts_dir, "fewshot_cot", "fewshot_cot_prompt_rephrased.txt"
            )
            args.fewshot_ost_prompt_rephrased_path = os.path.join(
                prompts_dir, "fewshot_ost", "fewshot_ost_prompt_rephrased.txt"
            )
            args.decompose_prompt_rephrased_path = os.path.join(
                prompts_dir, "decompose", "decompose_prompt_rephrased.txt"
            )
        else:
            # 使用标准提示词作为重述版本
            args.fewshot_cot_prompt_rephrased_path = os.path.join(
                prompts_dir, "fewshot_cot", "fewshot_cot_prompt.txt"
            )
            args.fewshot_ost_prompt_rephrased_path = os.path.join(
                prompts_dir, "fewshot_ost", "fewshot_ost_prompt.txt"
            )
            args.decompose_prompt_rephrased_path = os.path.join(
                prompts_dir, "decompose", "decompose_prompt.txt"
            )

    # ================================ 最终配置与启动 ================================
    args = post_process_args(args)  # 后处理参数，创建输出目录等
    print(args)                     # 打印完整参数配置用于日志记录
    save_args(args)                 # 保存参数配置到文件
    main(args)                      # 启动主程序
